import argparse
import logging
import os
from contextlib import asynccontextmanager
from pathlib import Path

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
from starlette.status import HTTP_200_OK

# 解析命令行参数
parser = argparse.ArgumentParser(description='Start the FastAPI .')
parser.add_argument('--env', default='development', help='Environment to run the  in')
args = parser.parse_args()

# 动态选择 .env 文件
env_file = f'.env.{args.env}'
load_dotenv(dotenv_path=Path('') / env_file)

from app.routers.qq import qq  # 导入qq.py中的路由函数
from app.routers.qq1 import qq1  # 导入qq.py中的路由函数
from app.models import http_resp

# 后端服务地址
spring_boot_address = os.getenv('spring_boot_address')
# comfy ui api地址
server_address = os.getenv('server_address')
# # comfy ui socket地址
websocket_address = os.getenv('websocket_address')
# OSS 配置
access_key_id = os.getenv('access_key_id')
access_key_secret = os.getenv('access_key_secret')
bucket_name = os.getenv('bucket_name')
endpoint = os.getenv('endpoint')

# mq配置
rabbit_mq_host = os.getenv('rabbit_mq_host')
rabbit_mq_port = os.getenv('rabbit_mq_port')
rabbit_mq_user = os.getenv('rabbit_mq_user')
rabbit_mq_password = os.getenv('rabbit_mq_password')
rabbit_mq_virtual_host = os.getenv('rabbit_mq_virtual_host')
rabbit_mq_exchange = os.getenv('rabbit_mq_exchange')
queue_name = os.getenv('queue_name')
while_list = os.getenv('while_list')

log = logging.getLogger("uvicorn")


async def catch_exceptions_middleware(request: Request, call_next):
    token = request.headers.get('X-Token')
    log.info(f"token {token}")
    try:
        path = request.url.path
        log.info(f"path {path}")
        log.info(while_list)
        if path not in while_list:
            if not token:
                raise HTTPException(status_code=401, detail='X-Token header is missing')
            # 这里可以添加更复杂的鉴权逻辑，例如与数据库中的token比对等
            if token != "expected_token":
                raise HTTPException(status_code=401, detail='Invalid X-Token')
        return await call_next(request)
    except HTTPException as e:
        content = http_resp.response_fail(msg=e.detail, code=e.status_code).model_dump()
        log.info("HTTP异常 %s", content)
        return JSONResponse(content=http_resp.response_fail(msg=e.detail, code=e.status_code).model_dump(),
                            status_code=HTTP_200_OK)
    except Exception as e:
        # log.exception("服务异常")
        content = http_resp.response_fail(msg="服务异常", data=str(e)).model_dump()
        log.info("服务异常：%s", content)
        return JSONResponse(content=content,
                            status_code=HTTP_200_OK)


# 定义一个异步上下文管理器来处理启动和关闭事件
@asynccontextmanager
async def lifespan(fastApi: FastAPI):
    # 启动时执行的代码
    log.info("Service is starting up...")
    # 在应用启动时执行的代码
    log.info(f"后端服务地址：{spring_boot_address}")
    log.info(f"comfy ui api地址：{server_address}")
    log.info(f"comfy ui socket地址：{websocket_address}")
    log.info(f"mq地址：{rabbit_mq_host}")
    yield
    # 关闭时执行的代码
    log.info("Service is shutting down...")


# 创建FastAPI实例时传入lifespan参数
fastApi = FastAPI(lifespan=lifespan)

# 全局异常处理中间件
fastApi.middleware('http')(catch_exceptions_middleware)

# 处理其他api接口
fastApi.include_router(qq, prefix='/api/v1')
fastApi.include_router(qq1, prefix='/api/v2')

if __name__ == "__main__":
    uvicorn.run(fastApi, host="0.0.0.0", port=8080, log_config='./uvicorn_config.json')
